package tcp_server

import (
	"context"
	"errors"
	"fmt"
	"net"
	"runtime"
	"sync/atomic"
	"time"

	"github.com/youkelike/gateway/core/log"
)

var (
	ErrServerClosed     = errors.New("tcp: Server closed")
	ErrAbortHandler     = errors.New("tcp: abort TCPHandler")
	LocalAddrContextKey = &contextKey{"local-addr"}
)

type contextKey struct {
	name string
}

func (k *contextKey) String() string {
	return "tcp_proxy context value " + k.name
}

type TCPHandler interface {
	ServeTCP(ctx context.Context, conn net.Conn)
}

type Server struct {
	Addr    string
	Handler TCPHandler
	BaseCtx context.Context

	WriteTimeout     time.Duration
	ReadTimeout      time.Duration
	KeepAliveTimeout time.Duration

	inShutdown int32
	doneChan   chan struct{}
	l          *net.TCPListener
}

func (s *Server) shuttingDown() bool {
	return atomic.LoadInt32(&s.inShutdown) != 0
}

func (srv *Server) ListenAndServe() error {
	if srv.shuttingDown() {
		return ErrServerClosed
	}
	if srv.doneChan == nil {
		srv.doneChan = make(chan struct{})
	}

	addr, err := net.ResolveTCPAddr("tcp", srv.Addr)
	if err != nil {
		return errors.New("need addr")
	}
	ln, err := net.ListenTCP("tcp", addr)
	if err != nil {
		return err
	}
	return srv.Serve(ln)
}

func (srv *Server) Serve(l *net.TCPListener) error {
	srv.l = l
	defer srv.l.Close()

	if srv.BaseCtx == nil {
		srv.BaseCtx = context.Background()
	}

	for {
		rw, e := l.AcceptTCP()
		if e != nil {
			select {
			case <-srv.doneChan:
				return ErrServerClosed
			default:
			}
			fmt.Printf("accept fail, err: %v\n", e)
			continue
		}

		go srv.serve(srv.BaseCtx, rw)
	}
}

func (srv *Server) serve(ctx context.Context, rwc net.Conn) {
	// 设置参数
	if d := srv.ReadTimeout; d != 0 {
		rwc.SetReadDeadline(time.Now().Add(d))
	}
	if d := srv.WriteTimeout; d != 0 {
		rwc.SetWriteDeadline(time.Now().Add(d))
	}
	if d := srv.KeepAliveTimeout; d != 0 {
		// 操作系统保活消息的发送设置
		if tcpConn, ok := rwc.(*net.TCPConn); ok {
			tcpConn.SetKeepAlive(true)
			tcpConn.SetKeepAlivePeriod(d)
		}
	}

	remoteAddr := rwc.RemoteAddr().String()
	defer func() {
		if err := recover(); err != nil && err != ErrAbortHandler {
			const size = 64 << 10
			buf := make([]byte, size)
			buf = buf[:runtime.Stack(buf, false)]
			log.Error("tcp: panic serving %v: %v\n%s", remoteAddr, err, buf)
		}
		// 回到这里就关闭连接
		rwc.Close()
	}()

	ctx = context.WithValue(ctx, LocalAddrContextKey, rwc.LocalAddr())
	if srv.Handler == nil {
		log.Fatal("handler empty")
	}
	srv.Handler.ServeTCP(ctx, rwc)
}

func (srv *Server) Close() error {
	if atomic.CompareAndSwapInt32(&srv.inShutdown, 0, 1) {
		close(srv.doneChan) // 关闭channel
		srv.l.Close()       // 执行listener关闭
	}
	return nil
}

func ListenAndServe(addr string, handler TCPHandler) error {
	server := &Server{
		Addr:     addr,
		Handler:  handler,
		doneChan: make(chan struct{}),
	}
	return server.ListenAndServe()
}
